
class Rectangle(object):
    __slots__ = ['xmin', 'ymin', 'xmax', 'ymax']
    def __init__(self, xmin, ymin, xmax, ymax):
        self.xmin = xmin
        self.ymin = ymin
        self.xmax = xmax
        self.ymax = ymax

    def overlaps(self, rect):
        """Returns true if this rectangle intersects the given rect."""
        if self.xmin > rect.xmax or self.xmax < rect.xmin:
            return False
        if self.ymin > rect.ymax or self.ymax < rect.ymin:
            return False
        return True

    def intersect(self, rect):
        """Returns the intersection of this rectangle with given rect,
        or None if they don't intersect."""
        if not self.overlaps(rect):
            return None
        xmin = max(self.xmin, rect.xmin)
        ymin = max(self.ymin, rect.ymin)
        xmax = min(self.xmax, rect.xmax)
        ymax = min(self.ymax, rect.ymax)
        return Rectangle(xmin, ymin, xmax, ymax)

    def union(self, rect):
        """Returns a rectangle that encloses both this rectangle
        and the given rect."""
        xmin = min(self.xmin, rect.xmin)
        ymin = min(self.ymin, rect.ymin)
        xmax = max(self.xmax, rect.xmax)
        ymax = max(self.ymax, rect.ymax)
        return Rectangle(xmin, ymin, xmax, ymax)
        
    def center(self):
        """Returns center point of rectangle."""
        return (0.5*(self.xmin + self.xmax), 0.5*(self.ymin + self.ymax))

    def contains_point(self, x, y):
        """Returns True if given point is inside the rect."""
        return (x >= self.xmin and x <= self.xmax and y >= self.ymin and y <= self.ymax)
    

class Point(object):
    __slots__ = ['x', 'y']
    def __init__(self, x, y):
        self.x = x
        self.y = y

node_count = 0

class PointQuadTreeNode(object):
    """Quad tree data structure to hold point features."""

    __slots__ = ['rect', 'sw_quad', 'nw_quad', 'se_quad', 'ne_quad',
                 'contents', 'has_children', 'max_points_per_node', 'xmid', 'ymid']
    
    def __init__(self, rect, max_points_per_node=1):
        """rect = (xmin, ymin, xmax, ymax) is bounding rectangle covered by this node.
        max_points_per_node is the maximum number of points that can be contained in
        a node before the node will split into subquads and the contents redistributed."""
        self.rect = rect
        # child nodes
        self.ne_quad = None
        self.nw_quad = None
        self.se_quad = None
        self.sw_quad = None
        # features contained in this node
        self.contents = None
        self.has_children = False
        self.max_points_per_node = max_points_per_node  # AKA threshold
        # Calculate and cache center point of rectangle.
        self.xmid, self.ymid = self.rect.center()

        global node_count
        node_count += 1

    def insert_point(self, point):
        """Inserts the point either into this node or one of its child nodes.
        Returns True if point successfully inserted, and False if point
        lies outside of the node rectangle."""
        if not self.rect.contains_point(point.x, point.y):
            # Point is not inside the rect covered by this node.
            return False
        if not self.has_children:
            # No child nodes yet, so first try to store the item in this node
            # if there is enough room for it.
            if not self.contents:
                self.contents = [ point ]
                return True
            else:
                self.contents.append(point)
                # Check to see if we've exceeded the number of items
                # that we're allowed to store per node.
                if len(self.contents) > self.max_points_per_node:
                    for item in self.contents:
                        self.push_point_to_child_node(item)
                    self.contents = None
                    self.has_children = True
                return True
        else:
            # If there are child nodes, then pass point on to the
            # correct child.
            self.push_point_to_child_node(point)
            return True

    def push_point_to_child_node(self, point):
        """Insert the given point into one of the child nodes.
        Also creates the necessary child node if it doesn't exist yet."""
        if point.x < self.xmid:
            if point.y < self.ymid:
                if not self.sw_quad:
                    subrect = Rectangle(self.rect.xmin, self.rect.ymin, self.xmid, self.ymid)
                    self.sw_quad = PointQuadTreeNode(subrect, self.max_points_per_node)
                self.sw_quad.insert_point(point)
            else:
                if not self.nw_quad:
                    subrect = Rectangle(self.rect.xmin, self.ymid, self.xmid, self.rect.ymax)
                    self.nw_quad = PointQuadTreeNode(subrect, self.max_points_per_node)
                self.nw_quad.insert_point(point)
        else:
            if point.y < self.ymid:
                if not self.se_quad:
                    subrect = Rectangle(self.xmid, self.rect.ymin, self.rect.xmax, self.ymid)
                    self.se_quad = PointQuadTreeNode(subrect, self.max_points_per_node)
                self.se_quad.insert_point(point)
            else:
                if not self.ne_quad:
                    subrect = Rectangle(self.xmid, self.ymid, self.rect.xmax, self.rect.ymax)
                    self.ne_quad = PointQuadTreeNode(subrect, self.max_points_per_node)
                self.ne_quad.insert_point(point)
        
    def get_points_in_rect(self, rect):
        """Return a list of all points contained in the rect."""
        # First check if rect intersects bounding rectangle for the node.
        if not self.rect.overlaps(rect):
            return []
        results = set()
        # If this node contains features itself, then run a per-feature query
        if self.contents:
            results.update([point for point in self.contents if rect.contains_point(point.x, point.y)])
        # If this node has children, run query on each child.
        if self.sw_quad:
            results.update(self.sw_quad.get_points_in_rect(rect))            
        if self.nw_quad:
            results.update(self.nw_quad.get_points_in_rect(rect))            
        if self.se_quad:
            results.update(self.se_quad.get_points_in_rect(rect))            
        if self.ne_quad:
            results.update(self.ne_quad.get_points_in_rect(rect))            

        return results

    def depth(self):
        depths = [0]
        if self.sw_quad: depths.append(self.sw_quad.depth())
        if self.nw_quad: depths.append(self.nw_quad.depth())
        if self.se_quad: depths.append(self.se_quad.depth())
        if self.ne_quad: depths.append(self.ne_quad.depth())                
        return 1 + max(depths)
                
class PointQuadTree(object):
    """Wrapper class for root PointQuadTreeNode."""
    def __init__(self, rect, max_points_per_node=1):
        self.root = PointQuadTreeNode(rect, max_points_per_node)
        
    def insert_point(self, point):
        return self.root.insert_point(point)

    def get_points_in_rect(self, rect):
        return self.root.get_points_in_rect(rect)

    def depth(self):
        return self.root.depth()


def create_test_data(n):
    import random
    points = []
    print 'Creating random points'
    for i in xrange(n):
        x = random.uniform(0,100)
        y = random.uniform(0,100)
        points.append( Point(x,y) )
    #f = open('testdata.py', 'w')
    #f.write('testdata = %s' % str(points))
    #f.close()
    print 'Building first quadtree'
    q1 = get_quadtree(points, 25)
    print 'Building second quadtree'
    q2 = get_quadtree(points, n+1)
    return q1, q2

def minimum_bounding_rectangle(points):
    x,y = points[0].x, points[0].y
    rect = Rectangle(x,y,x,y)
    for p in points:
        if p.x < rect.xmin:
            rect.xmin = p.x
        elif p.x > rect.xmax:
            rect.xmax = p.x
        if p.y < rect.ymin:
            rect.ymin = p.y
        elif p.y > rect.ymax:
            rect.ymax = p.y
    return rect
        
def get_quadtree(points, maxpoints=1):
    rect = minimum_bounding_rectangle(points)
    print 'done finding bounding rect'
    q = PointQuadTree(rect, maxpoints)
    for p in points:
        q.insert_point(p)
    return q
    

r = Rectangle(10,10,15,15)


## NOTES:

# 1. rewrite for rect features
# 2. don't create all child nodes when subdividing
#    only create those that are used.  (to save mem)
#    but will have to be able to eventually create if necessary.
# 3. for point queries, if node extent is contained inside the query rect
#    then should just return all contents without further testing.

######################################################################

class RectQuadTreeNode(object):
    """Quad tree node to hold rect features."""
    
    def __init__(self, rect, max_items_per_node=1):
        """rect defines the extent covered by this node.
        max_features_per_node is number of features that this
        node may contain before it is subdivided."""
        # Rectangular extent covered by node.
        self.rect = rect  
        # Child nodes
        self.ne_quad = None
        self.nw_quad = None
        self.se_quad = None
        self.sw_quad = None
        self.has_children = False
        self.contents = None
        self.max_items_per_node = max_items_per_node
        # Cache center point of rect extent.
        self.xmid, self.ymid = rect.center()
        
    def insert_item(self, item):
        """Inserts item into this node or one of it's children.
        item must have a rect attribute which is a minimum bounding
        rectangle around all features in the item.
        Returns True if item successfully inserted."""
        if not self.rect.overlaps(item.rect):
            # no part of item is inside this rect
            return False
        if not self.has_children:
            # node has no children yet, so try to store item in this node.
            if self.contents == None: self.contents = []
            self.contents.append(item)
            if len(self.contents) > self.max_items_per_node:
                # exceeded max number of items per node, so
                # we subdivide and redistribute contents to child nodes.
                items = self.contents
                self.contents = None
                for x in items:
                    self.push_item_to_child_node(x)
                self.has_children = True
            return True
        else:
            # node has children, so pass item on to correct child.
            self.push_item_to_child_node(x)
            return True

    def push_item_to_child_node(self, item):
        # Can assume that item rect overlaps the node rect, so that
        # item must overlap at least one of the child node rects.
        in_se = in_ne = in_sw = in_nw = False
        if item.rect.xmin < self.xmid:  
            if item.rect.ymin < self.ymid: in_sw = True
            if item.rect.ymax >= self.ymid: in_nw = True
        if item.rect.xmax >= self.xmid:  
            if item.rect.ymin < self.ymid: in_se = True
            if item.rect.ymax >= self.ymid: in_ne = True

        if in_sw:
            if not self.sw_quad:
                subrect = Rectangle(self.rect.xmin, self.rect.ymin, self.xmid, self.ymid)
                self.sw_quad = RectQuadTreeNode(subrect, self.max_items_per_node)
            self.sw_quad.insert_item(item)
        if in_nw:
            if not self.nw_quad:
                subrect = Rectangle(self.rect.xmin, self.ymid, self.xmid, self.rect.ymax)
                self.nw_quad = RectQuadTreeNode(subrect, self.max_items_per_node)
            self.nw_quad.insert_item(item)
        if in_se:
            if not self.se_quad:
                subrect = Rectangle(self.xmid, self.rect.ymin, self.rect.xmax, self.ymid)
                self.se_quad = RectQuadTreeNode(subrect, self.max_items_per_node)
            self.se_quad.insert_item(item)
        if in_ne:
            if not self.ne_quad:
                subrect = Rectangle(self.xmid, self.ymid, self.rect.xmax, self.rect.ymax)
                self.ne_quad = RectQuadTreeNode(subrect, self.max_items_per_node)
            self.ne_quad.insert_item(item)

    def get_items_in_rect(self, rect):
        if not self.rect.overlaps(rect):
            return []
        results = set()
        if self.contents:
            results.update([item for item self.contents if rect.overlaps(item.rect)])
        if self.sw_quad:
            results.update(self.sw_quad.get_items_in_rect(rect))
        if self.nw_quad:
            results.update(self.nw_quad.get_items_in_rect(rect))
        if self.se_quad:
            results.update(self.se_quad.get_items_in_rect(rect))
        if self.se_quad:
            results.update(self.ne_quad.get_items_in_rect(rect))            

        return results

    def depth(self):
        depths = [0]
        if self.sw_quad: depths.append(self.sw_quad.depth())
        if self.nw_quad: depths.append(self.nw_quad.depth())
        if self.se_quad: depths.append(self.se_quad.depth())
        if self.ne_quad: depths.append(self.ne_quad.depth())                
        return 1 + max(depths)
